module TopRoutine

using TopolParser
using PdbParser
using GroParser
using Structure
using Utils
using Printf

function top_functions(top::Topol, topfilename::String)
    while true
        println(h("\n                    ************ Topology functions ************", fc="c"))
        println("-10 Return to main menu")
        println("-1 Show topology summary")
        println(" 0 Write molecule topology to a new file")
        println(" 1 Add new topology items")
        println(" 2 Modify topology items")
        println(" 3 Delete topology items")
        println(" 4 View topology items")
        println(" 5 Append another itp file")
        println(" 6 Rearrange topology atoms")
        println(" 7 Convert to rtp file and generate hdb file")
        i = sel_fun(Int32)
        if i == -10
            return
        elseif i == -1
            show_topol_summary(top)
        elseif i == 0
            out_path = get_path(topfilename)
            println("Generating itp of $(top.moleculetype)...")
            to_itp(top, out_path)
            finish_path(out_path)
        elseif i == 1
            add_topol_items!(top)
        elseif i == 2
            modify_topol_items!(top)
        elseif i == 3
            delete_topol_items!(top)
        elseif i == 4
            view_topol_items(top)
        elseif i == 5
            append_topol_items!(top)
        elseif i == 6
            sort_topol_atoms(top)
        elseif i == 7
            prev_atoms, next_atoms, prev_con_atom, next_con_atom = get_exclude_atoms(top)
            top2rtp(top, topfilename; prev_atoms, next_atoms, prev_con_atom, next_con_atom)
            top2hdb(top, topfilename; prev_atoms, next_atoms)
        else
            println("Coming soon.")
        end
    end
end

function show_topol_summary(top)
    title_hint("\nTopology of " * top.moleculetype * " contains:\n")
    println(length(top.atomtypes), " types of atoms,")
    println(length(top.atoms), " atoms,")
    println(length(top.bonds), " bonds,")
    println(length(top.constraints), " constraints,")
    println(length(top.pairs), " pairs,")
    println(length(top.angles), " angles,")
    println(length(top.dihedrals), " dihedrals,")
    println(length(top.exclusions), " exclusions.")
end

function add_topol_items!(top::Topol)
    while true
        println()
        println(" 0 Return")
        println(" 1 Add atom topol items")
        println(" 2 Add bond topol items")
        println(" 3 Add constraints topol items")
        println(" 4 Add pairs topol items")
        println(" 5 Add angles topol items")
        println(" 6 Add dihedrals topol items")
        println(" 7 Add exclusions topol items")
        input = parse(Int32, readline())
        if input == 1
            request_hint("Input new atom line, including (splitted by spaces): ")
            println("atom number, atom type, residue number, residue name, atom name, atom charge group number, atom charge, and atom mass (optional)")
            println("(e.g., 1   CT   1   ALA   CA   1   0.3   12)")
            atom_top = gen_topol_atom_from_line(get_input(""))
            println("input insert position of the target atom (default: at the end)")
            pos = get_input(-1, Int32)
            add_atom!(top.atoms, atom_top; pos)
            finish_hint("Finished adding atom topology item.\n")
        elseif input == 2
            request_hint("Input new bond line, including (splitted by spaces): ")
            println("bond atom numbers, bond function type, bond parameters\n(e.g., 1   2   1    0.34    1e5)")
            bond_top = gen_topol_bond_from_line(top.atoms, get_input(""))
            add_bond!(top.bonds, bond_top)
            finish_hint("Finished adding bond topology item.\n")
        elseif input == 3
            request_hint("Input new constraint line, including (splitted by spaces): ")
            println("constraint atom numbers, constraint function type, constraint parameters\n(e.g., 1   2   1   0.27)")
            constraint_top = gen_topol_constraint_from_line(top.atoms, get_input(""))
            add_constraint!(top.constraints, constraint_top)
            finish_hint("Finished adding constraint topology item.\n")
        elseif input == 4
            request_hint("Input new pair line, including (splitted by spaces): ")
            println("pair atom numbers, pair function type\n(e.g., 1   2   1)")
            pair_top = gen_topol_pair_from_line(top.atoms, get_input(""))
            add_pair!(top.pairs, pair_top)
            finish_hint("Finished adding pair topology item.\n")
        elseif input == 5
            request_hint("Input new angle line, including (splitted by spaces): ")
            println("angle atom numbers, angle function type, angle parameters\n(e.g., 1   2   3   1   45   100)")
            angle_top = gen_topol_angle_from_line(top.atoms, get_input(""))
            add_angle!(top.angles, angle_top)
            finish_hint("Finished adding angle topology item.\n")
        elseif input == 6
            request_hint("Input new dihedral line, including (splitted by spaces): ")
            println("dihedral atom numbers, dihedral function type, dihedral parameters\n(e.g., 1   2   3   4   9   180   25   2)")
            dihedral_top = gen_topol_dihedral_from_line(top.atoms, get_input(""))
            add_dihedral!(top.dihedrals, dihedral_top)
            finish_hint("Finished adding dihedral topology item.\n")
        elseif input == 7
            request_hint("Input new exclusion line, including (splitted by spaces): ")
            println("exclusion atom numbers\n(e.g., 1 2 3 4 5)")
            exclusion_top = gen_topol_exclustion_from_line(top.atoms, get_input(""))
            add_exclusion!(top.exclusions, exclusion_top)
            finish_hint("Finished adding exclusion topology item.\n")
        elseif input == 0
            return
        end
        show_topol_summary(top)
    end
end

function modify_topol_items!(top::Topol)
    while true
        println()
        println("-1 Preview atoms")
        println(" 0 Return")
        println(" 1 Modify topol name")
        println(" 2 Shift atoms indexes")
        println(" 3 Shift atoms residue indexes")
        println(" 4 Multiply bonds force constant with a number")
        println(" 5 Multiply angles force constant with a number")
        println(" 6 Multiply dihedrals force constant with a number")
        para = sel_fun(Int32)
        if para == 0
            return
        elseif para == -1
            for atom in top.atoms
                println(atom)
            end
        elseif para == 1
            request_hint("Input the new molecule name (default: $(top.moleculetype)): ")
            top.moleculetype = get_input(top.moleculetype)
            finish_hint("Changing topol name finished: $(top.moleculetype)\n")
        elseif para == 2
            request_hint("Shift atoms indexes with (default: 0): ")
            shift_value = get_input(0, Int32)
            for atom in top.atoms
                atom.nr += shift_value
                atom.cgnr += shift_value
            end
            finish_hint("Shifted atoms indexes by $shift_value\n")
        elseif para == 3
            request_hint("Shift atoms indexes with (default: 0): ")
            shift_value = get_input(0, Int32)
            for atom in top.atoms
                atom.resnr += shift_value
            end
            finish_hint("Shifted atoms residue indexes by $shift_value\n")
        elseif para == 4
            request_hint("Multiply bonds force constant with (default: 1): ")
            multiply_value = get_input(1, Float64)
            for bond in top.bonds
                bond.c1 *= multiply_value
            end
            finish_hint("Multiplied bonds force constants by $multiply_value\n")
        elseif para == 5
            request_hint("Multiply angles force constant with (default: 1): ")
            multiply_value = get_input(1, Float64)
            for angle in top.angles
                angle.c1 *= multiply_value
            end
            finish_hint("Multiplied angles force constants by $multiply_value\n")
        elseif para == 6
            request_hint("Multiply dihedrals force constant with (default: 1): ")
            multiply_value = get_input(1, Float64)
            for dihedral in top.dihedrals
                dihedral.c1 *= multiply_value
            end
            finish_hint("Multiplied dihedrals force constants by $multiply_value\n")
        end
    end
end

function delete_topol_items!(top::Topol)
    while true
        println()
        println("-1 Show topology summary")
        println(" 0 Return")
        println(" 1 Delete atoms and related items (bonds, angles, etc.)")
        para = sel_fun(Int32)
        if para == 0
            return
        elseif para == -1
            show_topol_summary(top)
        elseif para == 1
            request_hint("Input atoms range to delete, e.g., 1-3, 5: ")
            atrange = atrange2atlist(get_input(""))
            println("Hard delete or soft delete?")
            println("Hard delete: delete all bonds, angles, dihedrals containing the selected atoms.")
            println("Soft delete: delete the bonds, angles, dihedrals consisting entirely of the selected atoms.")
            println(" 1 Hard delete")
            println(" 2 Soft delete")
            dm = sel_fun(Int32)
            if dm == 1
                filter!(b -> !(b.ai.nr in atrange || b.aj.nr in atrange), top.bonds)
                filter!(p -> !(p.ai.nr in atrange || p.aj.nr in atrange), top.pairs)
                filter!(c -> !(c.ai.nr in atrange || c.aj.nr in atrange), top.constraints)
                filter!(a -> !(a.ai.nr in atrange || a.aj.nr in atrange || a.ak.nr in atrange), top.angles)
                filter!(d -> !(d.ai.nr in atrange || d.aj.nr in atrange || d.ak.nr in atrange || d.al.nr in atrange), top.dihedrals)
                filter!(e -> !(any[(atom.nr in atrange) for atom in e.atoms]), top.exclusions)
                filter!(a -> !(a.nr in atrange), top.atoms)
                resort!(top.atoms)
            elseif dm == 2
                filter!(b -> !(b.ai.nr in atrange && b.aj.nr in atrange), top.bonds)
                filter!(p -> !(p.ai.nr in atrange && p.aj.nr in atrange), top.pairs)
                filter!(c -> !(c.ai.nr in atrange && c.aj.nr in atrange), top.constraints)
                filter!(a -> !(a.ai.nr in atrange && a.aj.nr in atrange && a.ak.nr in atrange), top.angles)
                filter!(d -> !(d.ai.nr in atrange && d.aj.nr in atrange && d.ak.nr in atrange && d.al.nr in atrange), top.dihedrals)
                filter!(e -> !(all[(atom.nr in atrange) for atom in e.atoms]), top.exclusions)
                filter!(a -> !(a.nr in atrange), top.atoms)
                resort!(top.atoms)
            else
                error_hint("Error input.\n")
            end
        end
        show_topol_summary(top)
    end
end

function view_topol_items(top::Topol)
    while true
        println()
        println(" 0 Return")
        println(" 1 Show topol atom types")
        println(" 2 Show topol atoms")
        println(" 3 Show topol bonds")
        println(" 4 Show topol constraints")
        println(" 5 Show topol pairs")
        println(" 6 Show topol angles")
        println(" 7 Show topol dihedrals")
        println(" 8 Show topol exclusions")
        para = sel_fun(Int32)
        if para == 0
            return
        elseif para == 1
            println("Atom types:")
            (x -> println(x)).(top.atomtypes)
        elseif para == 2
            println("Atoms:")
            (x -> println(x)).(top.atoms)
        elseif para == 3
            println("Bonds:")
            (x -> println(x)).(top.bonds)
        elseif para == 4
            println("Constraints:")
            (x -> println(x)).(top.constraints)
        elseif para == 5
            println("Pairs:")
            (x -> println(x)).(top.pairs)
        elseif para == 6
            println("Angles:")
            (x -> println(x)).(top.angles)
        elseif para == 7
            println("Dihedrals:")
            (x -> println(x)).(top.dihedrals)
        elseif para == 8
            println("Exclusions:")
            (x -> println(x)).(top.exclusions)
        end
    end
end

function append_topol_items!(top::Topol)
    request_hint("\nInput (absolute or relative) path of another itp file:\n")
    another_topfile = readline()
    if !isfile(another_topfile)
        println("Not a valid file.")
        return
    elseif !endswith(another_topfile, ".itp")
        println("Not an itp file.")
        return
    end
    another_top = gen_topol_from_file(another_topfile)
    request_hint("\nInput (absolute or relative) path of overlap file (default: no overlap, directly append):\n")
    title_hint("INFO: ")
    println("The overlap file (plain text, e.g., overlap.txt) includes four lines:\n",
            "1. main fragment of current residue (A1),\n",
            "2. extend fragment of next residue (A2),\n",
            "3. extend fragment of current residue (B1),\n",
            "4. main fragment of next residue (B2).\n",
            "Example:\n",
            "    4, 6, 9\n",
            "    1-3\n",
            "    1-3, 5, 7-8\n",
            "    4-7, 10, 12\n",
            "current residue: ... A1 —— A1 —— A1 -- B1 —— B1 —— B1 ...\n",
            "next residue:    ... A2 —— A2 —— A2 -- B2 —— B2 —— B2 ...\n",
            "$(h("NOTE", fc="y")): the atoms in the 1st/2nd line and 3rd/4th line must be in one-to-one correspondence!!!")
    overlap_file = readline()
    if !isempty(overlap_file) && !isfile(overlap_file)
        println("Not a valid file.")
        return
    else
        request_hint("Prefer preserving parameters in the current residue? [Y/n] ")
        current_first = get_input("Y")
        if current_first == "Y" || current_first == "y"
            current_first = true
        else
            current_first = false
        end
    end
    append_top!(top, another_top; overlap_file, current_first)
end

function sort_topol_atoms(top::Topol)
    println()
    println(" 0 Return")
    println(" 1 Sort by reference structure file")
    println(" 2 Sort by hand-input reference")
    input = sel_fun(Int32)
    if input == 0
        return
    elseif input == 1
        println("Input the reference structure file path:")
        struct_path = get_input("")
        atoms = Vector{GeneralAtom}()
        if endswith(struct_path, ".pdb")
            struct_file = PdbParser.gen_pdb_from_file(struct_path)
            atoms = struct_file.models[1].items
        elseif endswith(struct_path, ".gro")
            struct_file = GroParser.gen_gro_from_file(struct_path)
            atoms = struct_file.atoms
        else
            error_hint("Currently cannot handle this structure file.\n")
        end
        top.atoms = top.atoms[indexin(getfield.(atoms, :atname), getfield.(top.atoms, :atom))]
    elseif input == 2
        println("INFO: if you want the initial sequence [1 2 3 4 5] to be resorted into [1 4 5 2 3], \n",
                "then simply input \"1 4 5 2 3\" or \"1, 4-5, 2-3\" are both OK (no quotes).")
        println(" 0 Input directly.")
        println(" 1 Load from existing file.")
        m = sel_fun(Int32)
        if m == 0
            request_hint("Input the target sequence of topol items, separated by space:\n")
            tarseq = Utils.atrange2atlist(readline())
        elseif m == 1
            request_hint("Input the path of sequence file, e.g., E:/Data/new_top_sequence.txt (Note: file should contain only single line):\n")
            tarseq = Utils.atrange2atlist(readline(readline()))
        else
            error_hint("Error input.")
        end
        # 根据输入序列调整拓扑顺序
        top.atoms = (x -> top.atoms[x]).(tarseq)
        resort!(top.atoms)
        println("Sort topol atoms finished.")
    end
end

function get_exclude_atoms(top::Topol)
    request_hint("\nInput atoms id of the previous residue, e.g., 1-3, 5 (leave blank if it is the first residue): ")
    atlist = atrange2atlist(get_input(""))
    prev_atoms = Vector{TopolAtom}()
    try
        prev_atoms = top.atoms[atlist]
    catch BoundsError
        error_hint("Error: you selected atom indexs $(join(atlist, " ")), but $(top.moleculetype) only has $(length(top.atoms)) atoms.\n")
        return
    end
    if !isempty(prev_atoms)
        request_hint("Connection atom id of the previous residue (default: $(prev_atoms[1].nr)):\n")
        for atom in prev_atoms
            @printf("%3d:%8s\n", atom.nr, atom.atom)
        end
        prev_con_atom = top.atoms[get_input(prev_atoms[1].nr, Int32)]
        new_conname = startswith(prev_con_atom.atom, "C") ? "-C" : "-$(prev_con_atom.atom)"
        request_hint("Rename atom to? (default: $new_conname):\n")
        prev_con_atom.atom = get_input(new_conname)
    else
        prev_con_atom = nothing
    end
    request_hint("Input atoms id of the next residue, e.g., 2-4, 6 (leave blank if it is the last residue): ")
    atlist = atrange2atlist(get_input(""))
    next_atoms = Vector{TopolAtom}()
    try
        next_atoms = top.atoms[atlist]
    catch BoundsError
        error_hint("Error: you selected atom indexs $(join(atlist, " ")), but $(top.moleculetype) only has $(length(top.atoms)) atoms.\n")
        return
    end
    if !isempty(next_atoms)
        request_hint("Connection atom id of the next residue (default: $(next_atoms[1].nr)):\n")
        for atom in next_atoms
            @printf("%3d:%8s\n", atom.nr, atom.atom)
        end
        next_con_atom = top.atoms[get_input(next_atoms[1].nr, Int32)]
        new_conname = startswith(next_con_atom.atom, "N") ? "+N" : "+$(next_con_atom.atom)"
        request_hint("Rename atom to? (default: $new_conname):\n")
        next_con_atom.atom = get_input(new_conname)
    else
        next_con_atom = nothing
    end
    return prev_atoms, next_atoms, prev_con_atom, next_con_atom
end

function top2rtp(top::Topol, topfilename::String; prev_atoms=Vector{TopolAtom}(), next_atoms=Vector{TopolAtom}(), prev_con_atom=nothing, next_con_atom=nothing)
    request_hint("Input forcefield type, i.e., amber/gromos (default: amber): ")
    ff = get_input("amber")
    rtp_path = get_path(topfilename, false, "rtp")
    to_rtp(top, rtp_path, ff, prev_atoms, next_atoms, prev_con_atom, next_con_atom)
    finish_path(rtp_path)
end

struct HdbItem
    H_num::Int32
    H_type::Int32
    H_atom::AbstractString
    heavy_atoms::Vector{AbstractString}
end

function top2hdb(top::Topol, topfilename::String; prev_atoms, next_atoms)
    # atoms layout: H--i--j--k
    items = Vector{HdbItem}()
    # 1. 找到所有的H
    atoms_H = [atom for atom in top.atoms if startswith(atom.atom, "H") && !(atom in prev_atoms) && !(atom in next_atoms)]
    for atom_H in atoms_H
        # 2. 找到和H相连的原子i
        bond_i = filter(bond -> bond.ai.nr == atom_H.nr || bond.aj.nr == atom_H.nr, top.bonds)[1]   # H一定只连一个
        if bond_i.ai.nr == atom_H.nr
            atom_i = bond_i.aj
        else
            atom_i = bond_i.ai
        end
        # 3. 找到和i相连的原子j
        bonds_j = filter(bond -> bond.ai.nr == atom_i.nr || bond.aj.nr == atom_i.nr, top.bonds)
        atoms_j = Set{TopolAtom}()
        for bond in bonds_j
            if bond.ai.nr == atom_i.nr
                push!(atoms_j, bond.aj)
            else
                push!(atoms_j, bond.ai)
            end
        end
        atoms_j = [atoms_j...]
        atoms_j_heavy = filter(j -> !startswith(j.atom, "H"), atoms_j)
        # 4. 找到和j相连的原子k
        bonds_k = filter(bond -> bond.ai.nr == atoms_j_heavy[1].nr || bond.aj.nr == atoms_j_heavy[1].nr, top.bonds)
        atoms_k = Set{TopolAtom}()
        for bond in bonds_k
            if bond.ai.nr == atoms_j_heavy[1].nr
                push!(atoms_k, bond.aj)
            else
                push!(atoms_k, bond.ai)
            end
        end
        atoms_k = [atoms_k...]
        atoms_k_heavy = filter(k -> !startswith(k.atom, "H") && !(k.nr == atom_i.nr), atoms_k)      # 排除i
        # 5. 根据j中H和非H原子数判断H的类型
        num_noh = length(atoms_j_heavy)
        num_h = length(atoms_j) - num_noh
        if num_noh == 1
            if num_h == 1
                # type 2, 羟H
                push!(items, HdbItem(1, 2, atom_H.atom, [atom_i.atom, atoms_j_heavy[1].atom, atoms_k_heavy[1].atom]))
            elseif num_h == 2
                # type 3, 烯H/酰胺H
                if !(atom_H.atom[1:end-1] in getproperty.(items, :H_atom))
                    push!(items, HdbItem(2, 3, atom_H.atom[1:end-1], [atom_i.atom, atoms_j_heavy[1].atom, atoms_k_heavy[1].atom]))
                end
            elseif num_h == 3
                # type 4, 甲基H
                if !(atom_H.atom[1:end-1] in getproperty.(items, :H_atom))
                    push!(items, HdbItem(3, 4, atom_H.atom[1:end-1], [atom_i.atom, atoms_j_heavy[1].atom, atoms_k_heavy[1].atom]))
                end
            end
        elseif num_noh == 2
            if num_h == 1
                # type 1, 环H/肽H
                push!(items, HdbItem(1, 1, atom_H.atom, getproperty.([atom_i, atoms_j_heavy...], :atom)))
            elseif num_h == 2
                # type 6, 亚甲基H
                if !(atom_H.atom[1:end-1] in getproperty.(items, :H_atom))
                    push!(items, HdbItem(2, 6, atom_H.atom[1:end-1], getproperty.([atom_i, atoms_j_heavy...], :atom)))
                end
            end
        elseif num_noh == 3
            # type 5, 特丁基H
            push!(items, HdbItem(1, 5, atom_H.atom, getproperty.([atom_i, atoms_j_heavy...], :atom)))
        end
    end
    # 6. 输出文件
    hdb_path = get_path(topfilename, false, "hdb")
    hdb_file = open(hdb_path, "w")
    write(hdb_file, @sprintf("%-5s%d\n", top.moleculetype, length(items)))
    for item in items
        write(hdb_file, @sprintf("%-7d%-7d%-7s", item.H_num, item.H_type, item.H_atom))
        for ha in item.heavy_atoms
            write(hdb_file, @sprintf("%-7s", ha))
        end
        write(hdb_file, "\n")
    end
    close(hdb_file)
    finish_path(hdb_path)
end

end